'''
Created on Nov 15, 2010

@author: yev
'''
'''
Created on Sep 10, 2010

@author: yev
'''
import array
from mpi4py import MPI
import xml.etree.cElementTree as cET
import os
import Queue
from random import randint, choice
import string
import sys
import tempfile
import time
from threading import Thread
import World 
import collections

import numpy

private_queues = {}
results_queue =  collections.deque()

pipe_counter = 0

#EXTRA_SPECIAL_MESSAGE = "I ARE WEASEL"
new_line = "\n"

def str2bool(v):
    if v is None:
        return False
    return v.lower() in ["yes", "true", "t", "1"]


def find_ranks_assigned_to_task(rank_to_task, task_of_interest):
    print ("Looking for task %d in " %(int(task_of_interest)))
    print (rank_to_task)
    ranks = []
    for rank in rank_to_task:
        if rank_to_task[rank] == task_of_interest:
            ranks.append(rank)
    print("Found ranks:")
    print(ranks)
    return ranks

def work(xml_file):
    pipe_counter = 0
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    name = MPI.Get_processor_name()
    # sending id
    comm.send(name, dest=0, tag=World.WHO)
    rank_to_task = {}
    rank_to_task = comm.recv(rank_to_task, source=0, tag=World.ASSIGNMENT)
    if not rank_to_task:
        return
    
    assigned_task = rank_to_task[rank]
    
    
    print ("I am Luigi %d on %s. My assigned task is %s. My PID is %d" %(rank, name, assigned_task, os.getpid()))
    pipes = {}
    commands = []
    
    for event, elem in cET.iterparse(xml_file, events = ("start", "end")):
        if event == "end" and "task" in elem.tag:
            print elem.attrib["id"]
            if elem.attrib["id"] == assigned_task:
    
                for item in elem.getiterator("send"):
                    if item.text is None:
                        send = None
                        print ("Luigi %d was not assigned a send task" %(rank))
                    else:
                        destination = int(item.text)
                        offset = item.attrib["offset"]
                        receivers = find_ranks_assigned_to_task(rank_to_task, str(destination))

                        send = Send(int(assigned_task), offset, receivers)
    
                for item in elem.getiterator("receive"):
                    if item.text is None:
                        receive = None
                    else:
                        source = int(item.text)
                        senders = find_ranks_assigned_to_task(rank_to_task, str(source))
                        receive = Receive(source,senders)
    
                for pipe in elem.getiterator("pipe"):
                    id = pipe.attrib["id"]
                    incoming = str2bool(pipe.attrib["incoming"])
                    if incoming:
                        private_queues[pipe_counter] = Queue.Queue()
                        if id.startswith("$"):
                            pipe_path = id.strip("$")
                            pipe = Pipe(incoming, private_queue_index = pipe_counter, set_path = pipe_path)
                        else:
                            pipe = Pipe(incoming, private_queue_index = pipe_counter)
                        pipe_counter += 1
                    else:    
                        pipe = Pipe(incoming)
                    pipes[id] = pipe
    
                for command_desc in elem.getiterator("command"):
                    if command_desc.text is None:
                        continue
                    command_text = command_desc.text
                    command_text = command_text.split()
                    translated_command = ''
                    for word in command_text:
                        index = word.find('%')
                        if index >= 0:
                            word = word.strip('%')
                            word = pipes[word].getPath()
                        translated_command += word
                        translated_command += ' '
                    
                    out_hint = command_desc.attrib["out"]
                    index = out_hint.find('%')
                    if index >= 0:                      # internal pipe
                        out_hint = out_hint.strip('%')
#                        print out_hint
                        std_out = pipes[out_hint].getPath()
                    elif out_hint == "QUEUE":
                        std_out = "QUEUE"
                    elif out_hint == "NONE":
                        std_out = None
                    else:
                        std_out = out_hint
                    command = Command(translated_command, std_out)
                    commands.append(command)
                    
                elem.clear() 
        
    threads = []
    for command in commands:
        pipe_counter += 1
        thread = RunExternalCommand(command, pipe_counter)
        thread.setName(command.getCommand())
        thread.start()
        threads.append(thread)
    
    if send is None and receive is not None:
        print ("Luigi %d is in RECEIVING mode" %(rank))
        do_receive(receive, pipes, rank)
    elif send is not None and receive is None:
        print ("Luigi %d is in SENDING mode" %(rank))
        do_send(send, threads)
    elif send != None and receive != None:
        print ("Luigi %d is in RECEIVING AND SENDING mode" %(rank))
        do_receive_and_send(receive, pipes, rank, send, threads)
    
    for thread in threads:
        thread.join()


class Command:
    def __init__(self, command, out):
        self.command = command
        self.stdOut = out
    
    def getCommand(self):
        commandList = []
        for word in self.command.split():
            result = ""
            result += word
            commandList.append(result)
        return commandList
    
    def getStdOut(self):
        return self.stdOut

class Send:
    def __init__(self, destination, offset, receivers):
        self.destination = destination
        self.offset = offset
        self.receivers = receivers

class Receive:
    def __init__(self, source, senders):
        self.source = source
        self.senders = senders


class Internal_Pipe:
    def __init__(self, queue_index, path):
        self.queue_index = queue_index
        self.path = path
        self.worker = PipeWorker(self.path, self.queue_index)
        self.worker.daemon = True
        self.worker.start()
    
    def writeData(self, data):
        private_queues[self.queue_index].put(data)
    
    def closeFile(self):
        self.worker.done = True
        private_queues[self.queue_index].join()
        self.worker.fh.close()



class Pipe:
    def __init__(self, incoming, private_queue_index = None, set_path = None):
        self.incoming = incoming
        if set_path is not None:
            self.path = set_path
            open(self.path, 'w').close()
        else:    
            self.path = tempfile.gettempdir() + "/" + randomString()
            os.mkfifo(self.path)
        
        if self.incoming == True:
            self.private_queue_index = private_queue_index    
            self.worker = PipeWorker(self.path, self.private_queue_index)
            self.worker.daemon = True
            self.worker.start()
             
    def writeData(self, data):
        private_queues[self.private_queue_index].put(data)
        return private_queues[self.private_queue_index].qsize()
    
    def closeFile(self):
        if self.incoming == True:
            self.worker.done = True
            private_queues[self.private_queue_index].join()
        return
    
    def getQueueSize(self):
        return private_queues[self.private_queue_index].qsize()
    
    def getPath(self):
        return self.path
            
    def __str__(self):
        return ("Hello" + self.path)
 
class PipeWorker(Thread):
    def __init__(self, path, private_queue_index):
        Thread.__init__(self)
        self.pipePath = path
        self.number = 0
        self.done = False
        self.private_queue_index = private_queue_index
        self.fh = None
    def run(self):
        time.sleep(1)
        self.fh = open(self.pipePath, 'w')
        while True:
            if self.done is True and private_queues[self.private_queue_index].empty():
                break                
            
            if not  private_queues[self.private_queue_index].empty():
                work = private_queues[self.private_queue_index].get()
            else:
                time.sleep(1)
                continue
            
#            print work
            if isinstance(work, str):
                self.fh.write(work)
            else:
                self.fh.write(work.tostring())
            self.number += 1   
            private_queues[self.private_queue_index].task_done()
        self.fh.close()
                           
        
class RunExternalCommand(Thread):
    def __init__(self, externalCommand, pipe_counter):
        Thread.__init__(self)
        self.externalCommand = externalCommand
        self.pipe_counter = pipe_counter
        private_queues[self.pipe_counter] = Queue.Queue(0)
        self.command = ""
        for word in externalCommand.getCommand():
            self.command += word
            self.command += " "

        if self.externalCommand.getStdOut() is None:
            self.output = None
            self.proc = os.popen(self.command)
        elif self.externalCommand.getStdOut() == "QUEUE":
            self.output = "QUEUE"
            self.proc = os.popen(self.command)
#        elif self.externalCommand.getStdOut() == "PIPE":
#            self.output = "PIPE"
#            self.proc = os.popen(self.command)
        else:
            self.output = self.externalCommand.getStdOut()
            self.proc = os.popen(self.command)
            
    def run(self):
#        if self.output == "PIPE":
#            fh = open(self.externalCommand.getStdOut(), 'w')
#            while True:
#                line = self.proc.readline()
#                sys.stdout.flush()
#                fh.write(line)
#                fh.flush()
#                if not line:
#                    fh.close()
#                    print ("Thread is BREAKING: command=%s, capture=%s" %(self.command, self.externalCommand.getStdOut()))
#                    break
        
        if self.output == "QUEUE":
            while True:
                line = self.proc.readline()
                results_queue.appendleft(line)
                sys.stdout.flush()
                if not line:
#                    results_queue.appendleft(EXTRA_SPECIAL_MESSAGE)
                    try:
                        self.proc.close()
                    except IOError:
                        pass
                    print ("Thread is BREAKING: command=%s, capture=%s, qSize=%d" %(self.command, self.externalCommand.getStdOut(), len(results_queue)))
                    break
       
        elif  self.output == None:
            print ("Thread in DIRECT mode: command=%s, capture=%s" %(self.command, self.externalCommand.getStdOut()))
            exit_code = self.proc.close()
        
        else:
            fh = open(self.externalCommand.getStdOut(), 'w')
            while True:
                line = self.proc.readline()
                sys.stdout.flush()
                fh.write(line)
                if not line:
                    fh.close()
                    try:
                        self.proc.close()
                    except IOError:
                        pass
                    
                    print ("Thread is BREAKING: command=%s, capture=%s" %(self.command, self.externalCommand.getStdOut()))
                    break

def randomString():
    chars = string.ascii_letters
    return "".join(choice(chars) for x in range(randint(6, 6)))

def prepare_message(offset, threads):
    start_time = time.time()
    data = []
    retrieved = 0
    eof = False
    while retrieved < offset:
        try:
            bleeding_edge = results_queue.pop()
        except IndexError: # Give the thread a little time to catch up and fill the dequeue
            time.sleep(0.5)
            continue
#        if pattern.match(bleeding_edge) is not None:
        if we_are_done(threads, results_queue):
#        if not are_threads_alive(threads) and len(results_queue) ==0: # Thread exited AND all messages doled out
            eof = True
            break
        data.append(bleeding_edge)
        retrieved += 1
    message = "".join(data)
    print("Message prepared in %f seconds" %(time.time() - start_time))
    return [message, eof]
    
def do_send(send, threads):
    dummy = array.array('i', range(2))
    status = MPI.Status()
    comm = MPI.COMM_WORLD
#    new_line = '\n'
    offset = int(send.offset)
#    pattern = re.compile(EXTRA_SPECIAL_MESSAGE)
    req = None
    receivers = send.receivers
            
    while True:
        for rank in receivers:
            if not comm.Iprobe(source=rank, tag=MPI.ANY_TAG, status=status): # is there anyone waiting?
                continue
            message, eof = prepare_message(offset, threads)
            start_time = time.time()
            comm.Recv([dummy, 2, MPI.INT], source=status.source, tag=send.destination) # who is waiting?
            
            print ("Attempting to send %d to %d " %(len(message), status.source))
            if req is not None: # Did the previous recipient receive the work?
                req.Wait()
            req = comm.Isend([message, len(message), MPI.CHAR], dest=status.source, tag=World.NEW_WORK)
            print ("Successful sending to %d (%f)" %(status.source, time.time() - start_time))
            
            if eof:
                print "Sender detected EOF"
                receivers = send.receivers
                while len(receivers) > 0:
                    for rank in receivers:
                        if not comm.Iprobe(source=rank, tag=MPI.ANY_TAG, status=status):
                            continue
                        print ("Sending SEACREST_OUT to %d" %(rank))
                        comm.Send([dummy, 2, MPI.INT], dest=rank, tag=World.SEACREST_OUT)
                        comm.Recv([dummy, 2, MPI.INT], source=rank, tag=World.SEACREST_OUT_ACK)
                        print ("Received SEACREST_OUT ACK from %d" %(rank))
                        receivers.remove(rank)
                return
                
def do_receive(receive, pipe_pieces, luigi_rank):
    dummy = array.array('i', range(2))
    comm = MPI.COMM_WORLD
    status = MPI.Status()
    number_of_senders = len(receive.senders)
    print receive.senders
    
    data = numpy.empty(10, dtype = 'c')

    while number_of_senders > 0:
        for rank in receive.senders:
            if not comm.Iprobe(source=rank, tag=MPI.ANY_TAG, status=status):
                time.sleep(0.1)
                continue
            
            print ("Receiving something from rank %d" %(rank))
            
            if status.tag == World.NEW_WORK:
                
                data = numpy.empty(status.count, dtype = 'c')
#                if data.size < status.count:
#                    data.resize(status.count, refcheck = False)
                
                comm.Recv([data, status.count, MPI.CHAR], source=rank, tag=World.NEW_WORK)
                feed_incoming_pipes_wait(data, pipe_pieces, False)
#                print ("Receiving from %d" %(rank))
                
#                print("receive came back")
            elif status.tag == World.SEACREST_OUT:
                print ("Luigi %d is receiving that Seacrest is out from %d" %(luigi_rank, rank))
                comm.Recv([dummy, 2, MPI.INT], source=rank, tag=World.SEACREST_OUT)
                comm.Send([dummy, 2, MPI.INT], dest=rank, tag=World.SEACREST_OUT_ACK)
                receive.senders.remove(rank)
                print ("%d Remain" %(len(receive.senders)))
                number_of_senders -= 1
                if number_of_senders == 0:
                    print ("Number of Senders for Luigi %d = %d"%(luigi_rank, number_of_senders))
                    for pipe in pipe_pieces:
                        if  pipe_pieces[pipe].incoming == True:
                            print ("Luigi %d closing pipe %s" %(luigi_rank, pipe))
                            pipe_pieces[pipe].closeFile()
                            print ("Luigi %d closed pipe %s" %(luigi_rank, pipe))
                    return

def do_receive_and_send(receive, pipe_pieces, luigi_rank, send, threads):
    dummy = array.array('i', range(2))
    comm = MPI.COMM_WORLD
    luigi_name = MPI.Get_processor_name()
    status = MPI.Status()
    number_of_senders = len(receive.senders)
    offset = int(send.offset)
    new_line = "\n"
#    pattern = re.compile(EXTRA_SPECIAL_MESSAGE)
    print receive.senders
    counter = 0
    data = numpy.empty(10, dtype = 'c')
    wrap_up = False
        
    while True:
        if number_of_senders > 0:
            for rank in receive.senders:
                comm.Send([dummy, 2, MPI.INT], dest=rank, tag=receive.source)              
                while not comm.Iprobe(source=rank, tag=MPI.ANY_TAG, status=status):
                    time.sleep(0.1)

                if status.tag == World.NEW_WORK:
                    data = numpy.empty(status.count, dtype = 'c')
                    comm.Recv([data, status.count, MPI.CHAR], source=rank, tag=World.NEW_WORK)
                    feed_incoming_pipes_wait(data, pipe_pieces, True, luigi_name,luigi_rank)
                    counter += 1
                elif status.tag == World.SEACREST_OUT:
                    comm.Send([dummy, 2, MPI.INT], dest=rank, tag=World.SEACREST_OUT_ACK)
                    number_of_senders -= 1
                    receive.senders.remove(rank)
                    if number_of_senders == 0:
                        print ("Number of Senders for Promiscuous Luigi %d = %d"%(luigi_rank, number_of_senders))
                        for pipe in pipe_pieces:
                            if  pipe_pieces[pipe].incoming == True:
                                pipe_pieces[pipe].closeFile()
                        wrap_up = True
                        
        received_data = []
        if len(results_queue) > 0:
            time.sleep(0.25)
            retrieved = 0
            for rank in send.receivers:
                while len(results_queue) > 0:
                    bleeding_edge =  results_queue.pop()
#                    if pattern.match(bleeding_edge) is not None:
                    if we_are_done(threads, results_queue):
#                    if not are_threads_alive(threads) and len(results_queue) == 0:
                        message = "".join(received_data)
                        comm.Send([message, len(message), MPI.CHAR], dest=rank, tag=World.NEW_WORK)
                        for rank in send.receivers:
                            comm.Send([dummy, 2, MPI.INT], dest=rank, tag=World.SEACREST_OUT)
                            print ("Promiscuous Luigi %d Sending SEACREST_OUT to %d" %(luigi_rank, rank))
                            comm.Recv([dummy, 2, MPI.INT], source=rank, tag=World.SEACREST_OUT_ACK)
                            print ("Promiscuous Luigi %d  Receiving SEACREST_OUT ACK from %d" %(luigi_rank, rank))
                        return
                    received_data.append(bleeding_edge)
                    retrieved += 1
                    if len(results_queue) == 0 or retrieved > offset:
                        break
                message = "".join(received_data)
                comm.Send([message, len(message), MPI.CHAR], dest=rank, tag=World.NEW_WORK)
                
                received_data = ''
        
        if len(results_queue) == 0 and wrap_up is True and not are_threads_alive(threads):
            for rank in send.receivers:
                comm.Send([dummy, 2, MPI.INT], dest=rank, tag=World.SEACREST_OUT)
                comm.Recv([dummy, 2, MPI.INT], source=rank, tag=World.SEACREST_OUT_ACK)
            return

def are_threads_alive(threads):
    for thread in threads:
        if thread.is_alive():
            return True
    return False


def we_are_done(threads, results_queue):
    if not are_threads_alive(threads):
        if len(results_queue) == 0: # the call is blocking so best to avoid it
            return True
    return False

def feed_incoming_pipes_wait(data, pipes, sleep = False, name = "None", rank = -1):
    min_buffer_size = sys.maxint
    for key in pipes:
        if pipes[key].incoming is True:
            buffer_size = pipes[key].writeData(data)
            if buffer_size < min_buffer_size:
#                print ("Buffer size = %d" %(buffer_size))
                min_buffer_size = buffer_size
    max_sleep_cycles = 600
    sleep_cycle = 0
    if sleep is True:
        while (min_buffer_size > 0) and (sleep_cycle < max_sleep_cycles):
            min_buffer_size = sys.maxint
            for key in pipes:
                if pipes[key].incoming is True:
                    buffer_size = pipes[key].getQueueSize()
                    if buffer_size < min_buffer_size:
                        min_buffer_size = buffer_size

            time.sleep(0.05)
            sleep_cycle += 1

    
